#include "dec_pro.hh"

using namespace std;

dec_pro::dec_pro(string base_txt, string rms_txt, string input_file_name, string root_file_name, string save_path, bool is_write_txt, int asic_num, string encoding_list_filename, bool is_allow_dec_discon, double LSB) : base(asic_num), noise(asic_num), idx_wait_dec(2), chl_wait_dec(2), amp_wait_dec(2)
{
    this->asic_num = asic_num;
    this->save_path = save_path;
    this->root_filename = root_file_name;
    this->is_write_txt = is_write_txt;
    this->LSB = LSB;

    // 导入编码表
    int nchn = 128 * (asic_num / 2);
    dec_task_cg = new hit_dec(encoding_list_filename, 40, nchn, nchn);
    dec_task_cg->is_allow_discontinue = is_allow_dec_discon;

    if (dec_task_cg->is_list_get)
    {
        log_file << "Encoding lists " << encoding_list_filename << " get successfully" << endl;
    }
    else
    {
        log_file << "[Error]: Encoding lists " << encoding_list_filename << " get failed!" << endl;
        return;
    }

    // 导入基线
    if (!read_base(base_txt))
    {
        return;
    }
    if (!read_noise(rms_txt)) {
      return;
    }
    // 导入数据
    if (read_root(input_file_name))
    {
        init_root();
        deduce_base();
        data_split();
        data_dec();
        complete_adding();
    }
    else
    {
        return;
    }
}

dec_pro::~dec_pro()
{
    if (adas_dec_file != nullptr && adas_dec_file->IsOpen())
    {
        adas_dec_file->Close();
        delete adas_dec_tree;
        delete adas_dec_file;
    }
}

bool dec_pro::read_base(string base_txt)
{
    ifstream base_file(base_txt);
    string line;
    if (!base_file.is_open())
    {
        cout << cRED << "Failed to open file: " << base_txt << endl;
        return false;
    }
    while (getline(base_file, line))
    {
        istringstream iss(line);
        vector<string> line_split;
        string s;
        while (getline(iss, s, ' '))
        {
            line_split.push_back(s);
        }

        vector<double> line_split_double;
        for (int i = 0; i < line_split.size(); i++)
        {
            line_split_double.push_back(stod(line_split[i]));
        }

        for (int i = 0; i < line_split_double.size(); i++)
        {
            base[i].push_back(line_split_double[i]);
        }
    }
    base_file.close();
    cout << cBLUE << "Read base done!" << cRESET << endl;
    return true;
}

bool dec_pro::read_noise(string noise_txt) {
  ifstream noise_file(noise_txt);
  string line;
  if (!noise_file.is_open()) {
    cout << cRED << "Failed to open file: " << noise_txt << endl;
    return false;
  }
  while (getline(noise_file, line)) {
    istringstream iss(line);
    vector<string> line_split;
    string s;
    while (getline(iss, s, ' ')) {
      line_split.push_back(s);
    }

    vector<double> line_split_double;
    for (int i = 0; i < line_split.size(); i++) {
      line_split_double.push_back(stod(line_split[i]));
    }

    for (int i = 0; i < line_split_double.size(); i++) {
      noise[i].push_back(line_split_double[i]);
    }
  }
  noise_file.close();
  cout << cBLUE << "Read noise done!" << cRESET << endl;
  return true;
}

bool dec_pro::read_root(string input_file_name)
{
    TFile *input_file = new TFile(input_file_name.c_str());
    if (input_file->IsZombie() || !input_file->IsOpen())
    {
        cout << cRED << "File" << input_file_name << " open failed. Please check the file" << cRESET << endl;
        return false;
    }
    TTree *input_tree = (TTree *)input_file->Get("adas_initial_data");
    input_tree->SetBranchAddress("nhits", &anhits);
    input_tree->SetBranchAddress("triggerid", &triggerid);
    input_tree->SetBranchAddress("adas_hit_idx", &adas_hit_idx);
    input_tree->SetBranchAddress("adas_hit_chl", &adas_hit_chl);
    input_tree->SetBranchAddress("adas_hit_amp", &adas_hit_amp);

    int nentries = input_tree->GetEntries();

    for (int i = 0; i < nentries; i++)
    {
        input_tree->GetEntry(i);
        vector<int> temp_idx;
        vector<int> temp_chl;
        vector<int> temp_amp;

        for (int j = 0; j < anhits; j++)
        {
            if (adas_hit_chl[j] > 128)
            {
                continue;
            }
            // [todo] rms not used
            // if (abs((adas_hit_amp[j] -
            // base[adas_hit_idx[j]][adas_hit_chl[j]-1]) * LSB) > (5 *
            // noise[adas_hit_idx[j]][adas_hit_chl[j]-1]))
            // {
            //     temp_idx.push_back(adas_hit_idx[j]);
            //     temp_chl.push_back(adas_hit_chl[j] - 1);
            //     temp_amp.push_back(adas_hit_amp[j]);
            // }
            temp_idx.push_back(adas_hit_idx[j]);
            temp_chl.push_back(adas_hit_chl[j] - 1); // 扣除基线和解码时，通道号需从0开始
            temp_amp.push_back(adas_hit_amp[j]);
        }
        if (temp_idx.size() != 0)
        {
            trig.push_back(triggerid);
            idx.push_back(temp_idx);
            chl.push_back(temp_chl);
            amp.push_back(temp_amp);
        }
    }
    input_file->Close();
    cout << cBLUE << "Read root file done!" << cRESET << endl;
    cout << endl;
    return true;
}

void dec_pro::deduce_base()
{
    for (int i = 0; i < idx.size(); i++)
    {
        for (int j = 0; j < idx[i].size(); j++)
        {
            amp[i][j] = amp[i][j] - base[idx[i][j]][chl[i][j]];
        }
    }

    cout << cBLUE << "Data deduce base done!" << cRESET << endl;
    cout << endl;
}

void dec_pro::data_split()
{

    // 分为两个维度
    for (int i = 0; i < idx.size(); i++)
    {
        for (int j = 0; j < idx[i].size(); j++)
        {
            if (idx[i][j] < (asic_num / 2))
            {

                chl[i][j] = chl[i][j] + (idx[i][j] * 128);
                idx[i][j] = 0;
            }
            if (idx[i][j] >= (asic_num / 2))
            {

                chl[i][j] = chl[i][j] + ((idx[i][j] - (asic_num / 2)) * 128);
                idx[i][j] = 1;
            }
        }
    }

    for (int i = 0; i < idx.size(); i++)
    {
        vector<int> idx_0;
        vector<int> idx_1;
        vector<int> chl_0;
        vector<int> chl_1;
        vector<int> amp_0;
        vector<int> amp_1;

        for (int j = 0; j < idx[i].size(); j++)
        {
            if (idx[i][j] == 0)
            {
                idx_0.push_back(idx[i][j]);
                chl_0.push_back(chl[i][j]);
                amp_0.push_back(amp[i][j]);
            }
            else if (idx[i][j] == 1)
            {
                idx_1.push_back(idx[i][j]);
                chl_1.push_back(chl[i][j]);
                amp_1.push_back(amp[i][j]);
            }
        }
        idx_wait_dec[0].push_back(idx_0);
        idx_wait_dec[1].push_back(idx_1);
        chl_wait_dec[0].push_back(chl_0);
        chl_wait_dec[1].push_back(chl_1);
        amp_wait_dec[0].push_back(amp_0);
        amp_wait_dec[1].push_back(amp_1);
    }

    cout << cBLUE << "Data split X/Y direction done!" << cRESET << endl;
    cout << endl;
}

void dec_pro::data_dec()
{
    // 每次运行前清空nhits 和 cluser_number，以免后续程序判断错误
    nhits = 0;
    cluster_number = 0;
    // Clear the data stored in root for not confusing while in TBrowser()
    for (int i = 0; i < cMAX_DEC_HITS; i++)
    {
        hit_strips[i] = 0;
        hit_asic_chn[i] = 0;
        hit_amp[i] = 0;
        hit_time[i] = 0;
        hit_max_position[i] = 0;
        dimension_idx[i] = 0;
        cluster_size[i] = 0;
        cluster_holed_num[i] = 0;
    }

    int event_num = trig.size();

    for (int i = 0; i < event_num; i++)
    {
        nhits = 0;
        cluster_number = 0;
        bool is_qx = 0;
        bool is_qy = 0;
        int tem = 0;
        for (int j = 0; j < 2; j++)
        {
            int dec_strip_cnt = 0;
            if (j != 0)
            {
                for (auto &element : idx_wait_dec[j][i])
                {
                    element = 0;
                }
            }
            if (dec_task_cg->run_dec(chl_wait_dec[j][i]))
            {
                for (int k = 0; k < dec_task_cg->cluster_num.size(); k++)
                {
                    for (int z = 0; z < dec_task_cg->cluster_num[k]; z++)
                    {
                        hit_strips[nhits] = dec_task_cg->hit_strip[z + dec_strip_cnt];
                        hit_asic_chn[nhits] = chl_wait_dec[j][i][dec_task_cg->hit_seq[z + dec_strip_cnt]];
                        hit_amp[nhits] = amp_wait_dec[j][i][dec_task_cg->hit_seq[z + dec_strip_cnt]];
                        hit_time[nhits] = 0;
                        hit_max_position[nhits] = 0;
                        dimension_idx[nhits] = j;
                        nhits++;
                    }
                    cluster_size[cluster_number] = dec_task_cg->cluster_num[k];
                    cluster_holed_num[cluster_number] = dec_task_cg->cluster_holed[k];
                    cluster_number++;
                    dec_strip_cnt += dec_task_cg->cluster_num[k];
                }
            }
            if (j == 0)
            {
                if (nhits != 0)
                {
                    is_qx = 1;
                    tem = nhits;
                }
            }
            else
            {
                if (nhits != 0 && nhits > tem)
                {
                    is_qy = 1;
                }
            }
        }
        // x/y 两个维度至少一个维度有数据才保存
        if (nhits != 0 && cluster_number != 0 && (is_qx || is_qy))
        {
            trigger_id = trig[i];
            adas_dec_tree->Fill();
            if (is_write_txt_file)
            {
                write_txt();
            }
        }
    }

    cout << cBLUE << "Data decode done!" << cRESET << endl;
    cout << endl;
}

bool dec_pro::init_root()
{
    string filename = root_filename + ".root";
    adas_dec_file = new TFile(filename.c_str(), "recreate");
    adas_dec_tree = new TTree("asic_hit_data", "FEC Hit");

    adas_dec_tree->Branch("nhits", &nhits, "nhits/I");
    adas_dec_tree->Branch("trigger_id", &trigger_id, "trigger_id/I");
    adas_dec_tree->Branch("dimension_idx", dimension_idx, "dimension_idx[nhits]/I");
    adas_dec_tree->Branch("hit_strips", hit_strips, "hit_strips[nhits]/I");
    adas_dec_tree->Branch("hit_asic_chn", hit_asic_chn, "hit_asic_chn[nhits]/I");
    adas_dec_tree->Branch("hit_time", hit_time, "hit_time[nhits]/D");
    adas_dec_tree->Branch("hit_max_position", hit_max_position, "hit_max_position[nhits]/D");
    adas_dec_tree->Branch("hit_amp", hit_amp, "hit_amp[nhits]/I");
    adas_dec_tree->Branch("cluster_number", &cluster_number, "cluster_number/I");
    adas_dec_tree->Branch("cluster_holed_num", cluster_holed_num, "cluster_holed_num[cluster_number]/I");
    adas_dec_tree->Branch("cluster_size", cluster_size, "cluster_size[cluster_number]/I");

    cout << cCYAN << "TTree in TFile " << cMAGENTA << filename << cCYAN << " init done" << cRESET << endl;
    cout << cGREEN << filename << cRESET << endl;
    cout << endl;

    if (is_write_txt)
    {
        txt_filename = root_filename + ".txt";
        txt_file.open(txt_filename, ios::out | ios::trunc);
        if (txt_file.is_open())
        {
            cout << cCYAN << "Save hit position txt file to " << cMAGENTA << txt_filename << cRESET << endl;
            is_write_txt_file = true;
        }
        else
        {
            cout << cRED << "Cannot create file: " << txt_filename << cRESET << endl;
            is_write_txt_file = false;
        }
    }
    return true;
}

void dec_pro::write_txt()
{
    txt_file.open(txt_filename, ios::app);
    txt_file << "{" << endl;
    txt_file << "trigger_id: " << trigger_id << endl;
    txt_file << "nhits: " << nhits << endl;
    int i = 0;
    for (int k = 0; k < cluster_number; k++)
    {
        txt_file << "-------- Cluster " << k << "--------" << endl;
        int base = i;
        txt_file << "    cluster size:\t" << cluster_size[k] << endl;
        txt_file << "    dimension_idx:\t" << dimension_idx[base] << endl;
        txt_file << "    hit_strips:\t\t";
        for (int j = 0; j < cluster_size[k]; j++)
        {
            txt_file << hit_strips[j + base];
            if (j != cluster_size[k] - 1)
            {
                txt_file << ","
                         << "\t";
            }
            i++;
        }
        txt_file << endl;
        txt_file << "    hit_amp:\t";
        for (int j = 0; j < cluster_size[k]; j++)
        {
            txt_file << hit_amp[j + base];
            if (j != cluster_size[k] - 1)
            {
                txt_file << ","
                         << "\t";
            }
        }
        txt_file << endl;
    }
    txt_file << "}" << endl;
    txt_file.close();
}

void dec_pro::complete_adding()
{
    if (adas_dec_tree != nullptr)
    {
        adas_dec_file->cd();
        adas_dec_tree->Write();
        adas_dec_file->Flush();
        adas_dec_file->Close();
    }
    else
    {
        cout << "adas_dec_tree is not initialized!" << std::endl;
    }
}
